package geppetto.cat.programs;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentSymetrization;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.alignments.output.AlignerOutputLatex;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;

import java.io.IOException;
import java.io.PrintStream;
import java.io.UnsupportedEncodingException;
import java.util.Random;



/**
 * Builds a latex report of the output of a given model
 * 
 * @author javg
 * 
 */
public class BuildOutputReport extends BaseExp {
	String corpusDescription;

	int size, maxSentenceSize;

	String modelDir;

	String modelName;

	float treshAERF, treshAERB;

	float treshPostPTSF, treshPostPTSB;

	AlignmentsSet fViter, bViter, fPostAER, bPostAER, fPostPTS, bPostPTS,
			iviter, uviter, gviter, iAER, uAER, gAER, iPTS, uPTS, gPTS;

	AbstractModel mhmm, mhmmb;

	int maxSentenceOutputLen;

	String printStream;

	public BuildOutputReport(String[] args) {
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceSize = Integer.parseInt(args[2]);
		modelDir = args[3];
		modelName = args[4];
		maxSentenceOutputLen = Integer.parseInt(args[5]);
		printStream = args[6];
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("ModelDir " + modelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("---");
	}

	public void initializeModels(String modelName) {
		modelDir = modelDir + "/" + modelName + "/" + forwardCorpus.getName()
				+ "/" + size + "/model/";
		System.out.println("Loading Model");
		if (modelName.equalsIgnoreCase("baseline")) {
			mhmm = HMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
			mhmmb = HMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");
		} else if (modelName.equalsIgnoreCase("agreement")) {
			mhmm = AgreementHMM.loadModel(forwardCorpus, backwardCorpus,
					modelDir + "/MHMM/");
			mhmmb = ((AgreementHMM) mhmm).getBackwardModel();
		} else {

			System.out.println("Unknown Model: " + modelName);
			System.exit(1);
		}
		System.out.println("Finished Loading Model");
	}

	// Shoudl print latex headers and experiment headers
	public void printCorpusStatistics() {

	}

	public void getAlignment() {

		fViter = mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		bViter = mhmmb.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		iviter = AlignmentSymetrization.intersect(fViter, bViter);
		uviter = AlignmentSymetrization.union(fViter, bViter);
		gviter = AlignmentSymetrization.growDiagAux(fViter, bViter, iviter,
				uviter);
		treshAERF = mhmm.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
		treshAERB = mhmmb.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
		fPostAER = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
				treshAERF,false,false);
		bPostAER = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
				treshAERB,false,false);
		iAER = AlignmentSymetrization.intersect(fPostAER, bPostAER);
		uAER = AlignmentSymetrization.union(fPostAER, bPostAER);
		gAER = AlignmentSymetrization.growDiagAux(fPostAER, bPostAER, iAER,
				uAER);

		// sample ~200 sentences from the training corpus
		System.out.println("Forward" + mhmm);
		System.out.println("Backward" + mhmmb);
		Random r = new Random(-9119629771948792409L);
		int totNumSent = forwardCorpus
				.getNumSentences(BilingualCorpus.TRAIN_CORPUS);
		int[] sample = new int[200];
		for (int i = 0; i < sample.length; i++) {
			sample[i] = r.nextInt(totNumSent);
		}
		// count number of aligned points with viterbi
		int numAlignmentsF = mhmm.viterbiAlignments(sample,
				BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
		int numAlignmentsB = mhmmb.viterbiAlignments(sample,
				BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
		// tune posterior to get that.
		treshPostPTSF = mhmm.tuneThresholdNumPoints(numAlignmentsF, sample,
				BilingualCorpus.DEV_CORPUS,false);
		treshPostPTSB = mhmmb.tuneThresholdNumPoints(numAlignmentsB, sample,
				BilingualCorpus.DEV_CORPUS,false);
		System.out.println(treshAERB + " " + treshAERF);
		System.out.println(treshPostPTSB + " " + treshPostPTSF);
		fPostPTS = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
				treshPostPTSF,false,false);
		bPostPTS = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
				treshPostPTSB,false,false);
		iPTS = AlignmentSymetrization.intersect(fPostPTS, bPostPTS);
		uPTS = AlignmentSymetrization.union(fPostPTS, bPostPTS);
		gPTS = AlignmentSymetrization.growDiagAux(fPostPTS, bPostPTS, iPTS,
				uPTS);
	}

	public void printAlignment(PrintStream out, Alignment forward,
			Alignment backward, Alignment inter, Alignment grow,
			Alignment gold, BilingualCorpus corpus, String caption)
			throws UnsupportedEncodingException {
		AlignerOutputLatex.outputWithGoldPostLatexForwardBackwardInterGrow(
				forward, backward, inter, grow, gold, corpus, out, caption);
		printAlignmentStatistics(out, forward, backward, inter, grow, gold,
				gold.reverse(), "");
	}

	public void printAlignmentStatistics(PrintStream out, Alignment forward,
			Alignment backward, Alignment inter, Alignment grow,
			Alignment fGold, Alignment bGold, String caption) {
		Evaluation fRes = AlignmentEvaluator.evaluate(forward, fGold);
		Evaluation bRes = AlignmentEvaluator.evaluate(backward, bGold);
		Evaluation iRes = AlignmentEvaluator.evaluate(inter, fGold);
		Evaluation gRes = AlignmentEvaluator.evaluate(grow, fGold);

		out.println("\\begin{tabular}{|c|cccccccc|}");
		out.println("\\hline");
		out.println("&\\multicolumn{8}{|c|}{" + caption + "}\\\\");
		out.println("" + fRes.latexTableHeader());
		out.println("F Tresh=" + fRes.toLatexTableEntry());
		out.println("B Tresh=" + bRes.toLatexTableEntry());
		out.println("Inter " + iRes.toLatexTableEntry());
		out.println("grow " + gRes.toLatexTableEntry());
		out.println("\\hline");
		out.println("\\end{tabular}");
	}

	public void printGlobalAlignmentStatistics(PrintStream out,
			AlignmentsSet forward, AlignmentsSet backward, AlignmentsSet inter,
			AlignmentsSet grow, AlignmentsSet fGold, AlignmentsSet bGold,
			float fTresh, float bTresh, String caption) {
		Evaluation fRes = AlignmentEvaluator.evaluate(forward, fGold);
		Evaluation bRes = AlignmentEvaluator.evaluate(backward, bGold);
		Evaluation iRes = AlignmentEvaluator.evaluate(inter, fGold);
		Evaluation gRes = AlignmentEvaluator.evaluate(grow, fGold);
		out.println("\\hline");
		out.println("&\\multicolumn{8}{|c|}{" + caption + "}\\\\");
		out.println("" + fRes.latexTableHeader());
		out.println("F Tresh=" + fTresh + fRes.toLatexTableEntry());
		out.println("B Tresh=" + bTresh + bRes.toLatexTableEntry());
		out.println("Inter " + iRes.toLatexTableEntry());
		out.println("grow " + gRes.toLatexTableEntry());
		out.println("\\hline");
	}

	public int[][] getUnalignedPerSize(AlignmentsSet predicted) {
		int[] buckets = new int[5];
		int[] nrSentence = new int[5];
		int[][] res = new int[2][5];
		for (int i = 0; i < predicted.size(); i++) {
			Alignment al = predicted.get(i);
			int size = al.getForeignLenght();
			if (size < 10) {
				nrSentence[0]++;
				buckets[0] += al.getNumberNulls();
			} else if (size < 20) {
				nrSentence[1]++;
				buckets[1] += al.getNumberNulls();
			} else if (size < 30) {
				nrSentence[2]++;
				buckets[2] += al.getNumberNulls();
			} else if (size < 40) {
				nrSentence[3]++;
				buckets[3] += al.getNumberNulls();
			} else {
				nrSentence[4]++;
				buckets[4] += al.getNumberNulls();
			}
		}
		res[0] = buckets;
		res[1] = nrSentence;
		return res;
	}

	public String printUnalignedOne(AlignmentsSet predicted) {
		String res = "";
		int[][] aux = getUnalignedPerSize(predicted);
		int[] buckets = aux[0];
		int[] nrSentence = aux[1];
		for (int i = 0; i < buckets.length; i++) {
			res += "&" + buckets[i];
		}
		for (int i = 0; i < buckets.length; i++) {
			res += "&" + (buckets[i] * 1.0 / nrSentence[i]);
		}
		return res + "\\\\";
	}

	public String printNrSentences(AlignmentsSet predicted) {
		String res = "";
		int[][] aux = getUnalignedPerSize(predicted);
		int[] nrSentences = aux[1];
		for (int i = 0; i < nrSentences.length; i++) {
			res += "&" + nrSentences[i];
		}
		for (int i = 0; i < nrSentences.length; i++) {
			res += "&" + nrSentences[i];
		}
		return res + "\\\\";
	}

	public void printUnalignedPerSize(PrintStream out) {
		out.println("\\clearpage");
		out.println("\\begin{tabular}{|c|ccccc||ccccc|}");
		out.println("&10&20&30&40&rest&10&20&30&40&rest	\\\\ \n \\hline");
		out.println("Sentences" + printNrSentences(fViter));
		out.println("Fviter" + printUnalignedOne(fViter));
		out.println("Bviter" + printUnalignedOne(bViter));
		out.println("Iviter" + printUnalignedOne(iviter));
		out.println("Gviter" + printUnalignedOne(gviter));
		out.println("\\hline");
		out.println("FPTS" + printUnalignedOne(fPostPTS));
		out.println("BPTS" + printUnalignedOne(bPostPTS));
		out.println("IPTS" + printUnalignedOne(iPTS));
		out.println("GPTS" + printUnalignedOne(gPTS));
		out.println("\\hline");
		out.println("FAER" + printUnalignedOne(fPostAER));
		out.println("BAER" + printUnalignedOne(bPostAER));
		out.println("IAER" + printUnalignedOne(iAER));
		out.println("GAER" + printUnalignedOne(gAER));
		out.println("\\end{tabular}");

	}

	public void printAlignments(PrintStream out)
			throws UnsupportedEncodingException {
		int size = forwardCorpus.getNumSentences(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet fGold = forwardCorpus.getGold();
		AlignmentsSet bGold = backwardCorpus.getGold();
		out.println("\\begin{tabular}{|c|cccccccc|}");
		printGlobalAlignmentStatistics(out, fViter, bViter, iviter, gviter,
				fGold, bGold, 0, 0, "Viterbi");
		printGlobalAlignmentStatistics(out, fPostPTS, bPostPTS, iPTS, gPTS,
				fGold, bGold, treshPostPTSF, treshPostPTSB, "Post PTS");
		printGlobalAlignmentStatistics(out, fPostAER, bPostAER, iAER, gAER,
				fGold, bGold, treshAERF, treshAERB, "Post AER");
		out.println("\\end{tabular}");
		printUnalignedPerSize(out);
		out.println("\\clearpage\\pagebreak");
		for (int i = 0; i < size; i++) {
			int fLen = forwardCorpus.getForeignSentenceLength(i,
					BilingualCorpus.TEST_CORPUS);
			int sLen = forwardCorpus.getSourceSentenceLength(i,
					BilingualCorpus.TEST_CORPUS);
			AlignmentsSet gold = forwardCorpus.getGold();
			if (fLen < maxSentenceOutputLen && sLen < maxSentenceOutputLen) {
				printAlignment(out, fViter.get(i), bViter.get(i),
						iviter.get(i), gviter.get(i), gold.get(i),
						forwardCorpus, "Viterbi");
				out.println("\\clearpage\\pagebreak");
				printAlignment(out, fPostPTS.get(i), bPostPTS.get(i), iPTS
						.get(i), gPTS.get(i), gold.get(i), forwardCorpus,
						"Posterior PTS");
				out.println("\\clearpage\\pagebreak");
				printAlignment(out, fPostAER.get(i), bPostAER.get(i), iAER
						.get(i), gAER.get(i), gold.get(i), forwardCorpus,
						"Posterior AER");
				out.println("\\clearpage\\pagebreak");
				out.println("\n\n");
			}
		}

	}

	public void printAligmentBySentecenLen(PrintStream out) {
		AlignmentsSet[] sfViter = fViter.splitBySentenceLen();
		AlignmentsSet[] sbViter = bViter.splitBySentenceLen();
		AlignmentsSet[] siViter = iviter.splitBySentenceLen();
		AlignmentsSet[] sgViter = gviter.splitBySentenceLen();
		AlignmentsSet[] sfAER = fPostAER.splitBySentenceLen();
		AlignmentsSet[] sbAER = bPostAER.splitBySentenceLen();
		AlignmentsSet[] siAER = iAER.splitBySentenceLen();
		AlignmentsSet[] sgAER = gAER.splitBySentenceLen();
		AlignmentsSet[] sfPTS = fPostPTS.splitBySentenceLen();
		AlignmentsSet[] sbPTS = bPostPTS.splitBySentenceLen();
		AlignmentsSet[] siPTS = iPTS.splitBySentenceLen();
		AlignmentsSet[] sgPTS = gPTS.splitBySentenceLen();
		AlignmentsSet[] fGold = forwardCorpus.getGold().splitBySentenceLen();
		AlignmentsSet[] bGold = backwardCorpus.getGold().splitBySentenceLen();

		for (int i = 0; i < sfViter.length; i++) {
			out.println("\\begin{tabular}{|c|cccccccc|}");
			printGlobalAlignmentStatistics(out, sfViter[i], sbViter[i],
					siViter[i], sgViter[i], fGold[i], bGold[i], 0, 0,
					"Viterbi " + i);
			printGlobalAlignmentStatistics(out, sfPTS[i], sbPTS[i], siPTS[i],
					sgPTS[i], fGold[i], bGold[i], treshPostPTSF, treshPostPTSB,
					"Post PTS " + i);
			printGlobalAlignmentStatistics(out, sfAER[i], sbAER[i], siAER[i],
					sgAER[i], fGold[i], bGold[i], treshAERF, treshAERB,
					"Post AER " + i);
			out.println("\\end{tabular}");
			out.println("\\clearpage\\pagebreak");
		}

	}

	public void run() throws IOException {
		initializeCorpus(corpusDescription, size, maxSentenceSize);
		initializeModels(modelName);
		PrintStream out = new PrintStream(printStream);
		printBegin(out);
		getAlignment();
		printAligmentBySentecenLen(out);
		printAlignments(out);
		printEnd(out);

	}

	public static void main(String[] args) throws IOException {
		BuildOutputReport exp = new BuildOutputReport(args);
		exp.run();
	}

}
